#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Blender脚本：将顶点动画转换为骨骼动画
适用于FrankMocap生成的顶点动画数据
"""

import bpy
import bmesh
import mathutils
import numpy as np
from mathutils import Vector, Matrix, Quaternion
import pickle
import os

class VertexToBoneConverter:
    def __init__(self):
        self.scene = bpy.context.scene
        self.frame_start = 1
        self.frame_end = 250
        
        # SMPL/SMPLX骨骼定义（24个主要关节）
        self.smpl_joints = [
            'pelvis',       # 0
            'left_hip',     # 1 
            'right_hip',    # 2
            'spine1',       # 3
            'left_knee',    # 4
            'right_knee',   # 5
            'spine2',       # 6
            'left_ankle',   # 7
            'right_ankle',  # 8
            'spine3',       # 9
            'left_foot',    # 10
            'right_foot',   # 11
            'neck',         # 12
            'left_collar',  # 13
            'right_collar', # 14
            'head',         # 15
            'left_shoulder', # 16
            'right_shoulder', # 17
            'left_elbow',   # 18
            'right_elbow',  # 19
            'left_wrist',   # 20
            'right_wrist',  # 21
            'left_hand',    # 22
            'right_hand'    # 23
        ]
        
        # 骨骼层级关系
        self.bone_hierarchy = {
            'pelvis': None,
            'left_hip': 'pelvis',
            'right_hip': 'pelvis', 
            'spine1': 'pelvis',
            'left_knee': 'left_hip',
            'right_knee': 'right_hip',
            'spine2': 'spine1',
            'left_ankle': 'left_knee',
            'right_ankle': 'right_knee',
            'spine3': 'spine2',
            'left_foot': 'left_ankle',
            'right_foot': 'right_ankle',
            'neck': 'spine3',
            'left_collar': 'spine3',
            'right_collar': 'spine3',
            'head': 'neck',
            'left_shoulder': 'left_collar',
            'right_shoulder': 'right_collar',
            'left_elbow': 'left_shoulder',
            'right_elbow': 'right_shoulder',
            'left_wrist': 'left_elbow',
            'right_wrist': 'right_elbow',
            'left_hand': 'left_wrist',
            'right_hand': 'right_wrist'
        }

    def load_animation_data(self, pkl_path):
        """加载动画数据"""
        print(f"Loading animation data from: {pkl_path}")
        
        try:
            with open(pkl_path, 'rb') as f:
                data = pickle.load(f)
            
            # 检查数据格式
            if 'frames' in data:
                # frame_list格式
                frames = data['frames']
                self.frame_end = len(frames)
                
                # 提取顶点数据
                vertices_list = []
                for frame in frames:
                    if 'pred_output_list' in frame and frame['pred_output_list']:
                        pred_output = frame['pred_output_list'][0]
                        if 'pred_vertices_smpl' in pred_output:
                            vertices_list.append(pred_output['pred_vertices_smpl'])
                
                return np.array(vertices_list)
                
            elif 'vertices' in data:
                # animation_data格式
                self.frame_end = data.get('num_frames', 250)
                return data['vertices']
            
            else:
                print("Unknown data format")
                return None
                
        except Exception as e:
            print(f"Error loading data: {e}")
            return None

    def create_armature(self, mesh_obj):
        """创建骨骼结构"""
        print("Creating armature...")
        
        # 创建骨骼对象
        bpy.ops.object.armature_add(enter_editmode=True)
        armature_obj = bpy.context.active_object
        armature_obj.name = "SMPL_Armature"
        armature = armature_obj.data
        armature.name = "SMPL_Armature"
        
        # 清除默认骨骼
        bpy.ops.armature.select_all(action='SELECT')
        bpy.ops.armature.delete()
        
        # 估算关节位置（基于顶点组）
        joint_positions = self.estimate_joint_positions(mesh_obj)
        
        # 创建骨骼
        for i, joint_name in enumerate(self.smpl_joints):
            bone = armature.edit_bones.new(joint_name)
            
            if joint_name in joint_positions:
                head_pos = joint_positions[joint_name]
                bone.head = head_pos
                
                # 设置骨骼尾部位置
                if i < len(self.smpl_joints) - 1:
                    tail_offset = Vector((0, 0, 0.1))  # 默认长度
                    bone.tail = head_pos + tail_offset
                else:
                    bone.tail = head_pos + Vector((0, 0, 0.1))
            
            # 设置父子关系
            parent_name = self.bone_hierarchy.get(joint_name)
            if parent_name and parent_name in armature.edit_bones:
                bone.parent = armature.edit_bones[parent_name]
        
        # 退出编辑模式
        bpy.ops.object.mode_set(mode='OBJECT')
        
        return armature_obj

    def estimate_joint_positions(self, mesh_obj):
        """估算关节位置"""
        # 简化的关节位置估算（基于网格边界框）
        mesh = mesh_obj.data
        vertices = [mesh_obj.matrix_world @ v.co for v in mesh.vertices]
        
        # 计算边界框
        min_x = min(v.x for v in vertices)
        max_x = max(v.x for v in vertices)
        min_y = min(v.y for v in vertices)
        max_y = max(v.y for v in vertices)
        min_z = min(v.z for v in vertices)
        max_z = max(v.z for v in vertices)
        
        center_x = (min_x + max_x) / 2
        center_y = (min_y + max_y) / 2
        center_z = (min_z + max_z) / 2
        
        height = max_z - min_z
        
        # 估算主要关节位置
        joint_positions = {
            'pelvis': Vector((center_x, center_y, min_z + height * 0.5)),
            'spine1': Vector((center_x, center_y, min_z + height * 0.6)),
            'spine2': Vector((center_x, center_y, min_z + height * 0.7)),
            'spine3': Vector((center_x, center_y, min_z + height * 0.8)),
            'neck': Vector((center_x, center_y, min_z + height * 0.9)),
            'head': Vector((center_x, center_y, min_z + height * 0.95)),
            
            # 左腿
            'left_hip': Vector((min_x + (max_x - min_x) * 0.3, center_y, min_z + height * 0.5)),
            'left_knee': Vector((min_x + (max_x - min_x) * 0.3, center_y, min_z + height * 0.25)),
            'left_ankle': Vector((min_x + (max_x - min_x) * 0.3, center_y, min_z + height * 0.05)),
            'left_foot': Vector((min_x + (max_x - min_x) * 0.3, min_y + (max_y - min_y) * 0.8, min_z)),
            
            # 右腿
            'right_hip': Vector((max_x - (max_x - min_x) * 0.3, center_y, min_z + height * 0.5)),
            'right_knee': Vector((max_x - (max_x - min_x) * 0.3, center_y, min_z + height * 0.25)),
            'right_ankle': Vector((max_x - (max_x - min_x) * 0.3, center_y, min_z + height * 0.05)),
            'right_foot': Vector((max_x - (max_x - min_x) * 0.3, min_y + (max_y - min_y) * 0.8, min_z)),
            
            # 左臂
            'left_collar': Vector((min_x + (max_x - min_x) * 0.2, center_y, min_z + height * 0.8)),
            'left_shoulder': Vector((min_x + (max_x - min_x) * 0.1, center_y, min_z + height * 0.8)),
            'left_elbow': Vector((min_x, center_y, min_z + height * 0.6)),
            'left_wrist': Vector((min_x - (max_x - min_x) * 0.1, center_y, min_z + height * 0.6)),
            'left_hand': Vector((min_x - (max_x - min_x) * 0.15, center_y, min_z + height * 0.6)),
            
            # 右臂
            'right_collar': Vector((max_x - (max_x - min_x) * 0.2, center_y, min_z + height * 0.8)),
            'right_shoulder': Vector((max_x - (max_x - min_x) * 0.1, center_y, min_z + height * 0.8)),
            'right_elbow': Vector((max_x, center_y, min_z + height * 0.6)),
            'right_wrist': Vector((max_x + (max_x - min_x) * 0.1, center_y, min_z + height * 0.6)),
            'right_hand': Vector((max_x + (max_x - min_x) * 0.15, center_y, min_z + height * 0.6)),
        }
        
        return joint_positions

    def setup_armature_modifier(self, mesh_obj, armature_obj):
        """设置骨骼修改器"""
        print("Setting up armature modifier...")
        
        # 选择网格对象
        bpy.context.view_layer.objects.active = mesh_obj
        
        # 添加骨骼修改器
        armature_modifier = mesh_obj.modifiers.new(name="Armature", type='ARMATURE')
        armature_modifier.object = armature_obj
        
        # 自动权重绑定
        bpy.context.view_layer.objects.active = armature_obj
        bpy.ops.object.select_all(action='DESELECT')
        mesh_obj.select_set(True)
        armature_obj.select_set(True)
        bpy.context.view_layer.objects.active = armature_obj
        
        bpy.ops.object.parent_set(type='ARMATURE_AUTO')

    def create_bone_animation_from_vertices(self, armature_obj, vertices_data):
        """从顶点动画创建骨骼动画"""
        print("Creating bone animation from vertex data...")
        
        if vertices_data is None or len(vertices_data) == 0:
            print("No vertex data available")
            return
        
        num_frames = len(vertices_data)
        self.frame_end = num_frames
        
        # 设置动画范围
        self.scene.frame_start = self.frame_start
        self.scene.frame_end = self.frame_end
        
        # 选择骨骼对象
        bpy.context.view_layer.objects.active = armature_obj
        bpy.ops.object.mode_set(mode='POSE')
        
        # 为每个骨骼创建动画
        for frame_idx in range(num_frames):
            self.scene.frame_set(frame_idx + self.frame_start)
            
            # 分析当前帧的顶点变化
            current_vertices = vertices_data[frame_idx]
            
            # 计算根部移动（简化：使用顶点中心）
            if frame_idx == 0:
                self.initial_center = np.mean(current_vertices, axis=0)
            
            current_center = np.mean(current_vertices, axis=0)
            root_translation = current_center - self.initial_center
            
            # 应用到根骨骼
            if 'pelvis' in armature_obj.pose.bones:
                root_bone = armature_obj.pose.bones['pelvis']
                root_bone.location = Vector(root_translation)
                root_bone.keyframe_insert(data_path="location")
            
            # 对其他骨骼应用基本变换（这里可以根据需要扩展更复杂的算法）
            for bone_name in self.smpl_joints[1:]:  # 跳过根骨骼
                if bone_name in armature_obj.pose.bones:
                    bone = armature_obj.pose.bones[bone_name]
                    
                    # 这里可以实现更复杂的骨骼变换计算
                    # 简化版本：添加轻微的随机变化
                    if frame_idx > 0:
                        # 基于顶点变化计算骨骼旋转（简化算法）
                        rotation_factor = np.linalg.norm(root_translation) * 0.1
                        bone.rotation_quaternion = Quaternion((1, rotation_factor * 0.1, 0, 0))
                        bone.keyframe_insert(data_path="rotation_quaternion")
        
        # 退出姿态模式
        bpy.ops.object.mode_set(mode='OBJECT')
        print(f"Animation created for {num_frames} frames")

    def convert_vertex_animation_to_bone(self, mesh_obj, pkl_path):
        """主转换函数"""
        print("Starting vertex to bone animation conversion...")
        
        # 1. 加载动画数据
        vertices_data = self.load_animation_data(pkl_path)
        if vertices_data is None:
            print("Failed to load animation data")
            return None
        
        print(f"Loaded {len(vertices_data)} frames of vertex data")
        
        # 2. 创建骨骼
        armature_obj = self.create_armature(mesh_obj)
        
        # 3. 设置骨骼修改器
        self.setup_armature_modifier(mesh_obj, armature_obj)
        
        # 4. 创建骨骼动画
        self.create_bone_animation_from_vertices(armature_obj, vertices_data)
        
        print("Conversion completed!")
        return armature_obj

def main():
    """主函数 - 在Blender中运行"""
    
    # 检查是否有选中的网格对象
    if not bpy.context.active_object or bpy.context.active_object.type != 'MESH':
        print("Please select a mesh object")
        return
    
    mesh_obj = bpy.context.active_object
    print(f"Converting mesh: {mesh_obj.name}")
    
    # pkl文件路径（需要修改为你的文件路径）
    pkl_path = "D:/ai_project/frank-mo-cap/output_pkl/animation.pkl"
    
    if not os.path.exists(pkl_path):
        print(f"PKL file not found: {pkl_path}")
        print("Please update the pkl_path variable in the script")
        return
    
    # 创建转换器并执行转换
    converter = VertexToBoneConverter()
    armature_obj = converter.convert_vertex_animation_to_bone(mesh_obj, pkl_path)
    
    if armature_obj:
        print("✅ Conversion successful!")
        print("- Created armature with SMPL bone structure")
        print("- Generated bone animation from vertex data")
        print("- Ready to export as FBX with bone animation")
    else:
        print("❌ Conversion failed")

# 如果在Blender中运行
if __name__ == "__main__":
    main()